{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7f9d3587-a07f-4657-a57d-8d28619081d0",
   "metadata": {
    "id": "XIyP_0r6zuVc"
   },
   "source": [
    "<!-- Banner Image -->\n",
    "<img src=\"https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/landingpage/brevdevnotebooks.png\" width=\"100%\">\n",
    "\n",
    "<!-- Links -->\n",
    "<center>\n",
    "  <a href=\"https://console.brev.dev\" style=\"color: #06b6d4;\">Console</a> •\n",
    "  <a href=\"https://brev.dev\" style=\"color: #06b6d4;\">Docs</a> •\n",
    "  <a href=\"/\" style=\"color: #06b6d4;\">Templates</a> •\n",
    "  <a href=\"https://discord.gg/NVDyv7TUgJ\" style=\"color: #06b6d4;\">Discord</a>\n",
    "</center>\n",
    "\n",
    "# Run BioMistral 7B 🤙\n",
    "\n",
    "Welcome!\n",
    "\n",
    "In this notebook and tutorial, we will download & run the new [BioMistral 7B](https://huggingface.co/BioMistral/BioMistral-7B) (paper [here](https://arxiv.org/abs/2402.10373). \n",
    "\n",
    "If you would like a fine-tuning guide, join our [Discord](https://discord.gg/y9428NwTh3) or message me on [X](https://x.com/harperscarroll) to let me know!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c180393-2631-4e44-97bd-47b8cced2060",
   "metadata": {
    "id": "SSBw-KpkyRga"
   },
   "source": [
    "#### Before we begin: A note on OOM errors\n",
    "\n",
    "If you get an error like this: `OutOfMemoryError: CUDA out of memory`, tweak your parameters to make the model less computationally intensive. I will help guide you through that in this guide, and if you have any additional questions you can reach out on the [Discord channel](https://discord.gg/y9428NwTh3) or on [X](https://x.com/harperscarroll).\n",
    "\n",
    "To re-try after you tweak your parameters, open a Terminal ('Launcher' or '+' in the nav bar above -> Other -> Terminal) and run the command `nvidia-smi`. Then find the process ID `PID` under `Processes` and run the command `kill [PID]`. You will need to re-start your notebook from the beginning. (There may be a better way to do this... if so please do let me know!)\n",
    "\n",
    "### Let's begin!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae808e1c-21bc-4b00-ae44-bbfd7d1ed82e",
   "metadata": {
    "id": "hWI-uRLEyRgb"
   },
   "source": [
    "## 1. Get & Set Up a GPU\n",
    "\n",
    "I used a GPU and dev environment from [brev.dev](https://brev.dev). Click the badge below to get your preconfigured instance:\n",
    "\n",
    "[![](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/landingpage/brevdeploynavy.svg)](https://console.brev.dev/environment/new?instance=A10G:g5.xlarge&name=biomistral&file=https://github.com/brevdev/notebooks/raw/main/biomistral.ipynb&python=3.10&cuda=12.0.1)\n",
    "\n",
    "Once you've checked out your machine and landed in your instance page, select the specs you'd like (I used Python 3.10 and CUDA 12.0.1; these should be preconfigured for you if you use the badge above) and click the \"Build\" button to build your verb container. Give this a few minutes.\n",
    "\n",
    "A few minutes after your model has started Running, click the 'Notebook' button on the top right of your screen once it illuminates (you may need to refresh the screen). You will be taken to a Jupyter Lab environment, where you can upload this Notebook.\n",
    "\n",
    "Note: You can connect your cloud credits (AWS or GCP) by clicking \"Org: \" on the top right, and in the panel that slides over, click \"Connect AWS\" or \"Connect GCP\" under \"Connect your cloud\" and follow the instructions linked to attach your credentials."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "327a9d63-0dc6-49fa-ac93-550ff93ab4e6",
   "metadata": {},
   "source": [
    "Now, run the cell below (Shift + Enter when you've highlighted the cell) to import the required libraries onto the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72283f98-53ac-4111-abfe-f995c8e8500b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -q -U torch torchvision torchaudio transformers jupyter ipywidgets bitsandbytes accelerate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e51b1aa8-1602-4266-9af7-ff05ea060bc3",
   "metadata": {},
   "source": [
    "## 2. Set Up BioMistral\n",
    "Now that we've installed the necessary libraries, let's pull the BioMistral model from Hugging Face."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7347593-a776-4fd6-88bb-737e590fc2ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eedc36d012d8489c974d91ed8848dd1f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/14.5G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05b391005e7948bfa9fff33181b1f266",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/111 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16759fa804c349ad8a8ed3431e2e82d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/1.42k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13eb3739d45348bf8340ccf82e14e9b3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/493k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83dd052593ab4043b82978e100b7a27c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.80M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d7c248dc94f4d0490befc5ef98801cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/72.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "model_id = \"BioMistral/BioMistral-7B\"\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    trust_remote_code=True,\n",
    ")\n",
    "\n",
    "eval_tokenizer = AutoTokenizer.from_pretrained(model_id, add_bos_token=True, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e75ffa4-1d51-411b-b82a-3d7d54ccc98e",
   "metadata": {},
   "source": [
    "## 3. Run the Model! \n",
    "Replace `eval_prompt` with your prompt string below, and run the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7988e989-432c-4f98-8e1b-383fbfa1da3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best way to 3D print a model of the human heart is to use an image from a real patient. However, this requires that the patient undergoes cardiac magnetic resonance imaging (MRI) and provides consent for the data to be used in research. This process can be time consuming and costly. Therefore, we created a library of open-source images of the human heart that can be downloaded directly from our website (https://www.printheart.com/). These images\n"
     ]
    }
   ],
   "source": [
    "eval_prompt = \"The best way to \"\n",
    "model_input = eval_tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    print(eval_tokenizer.decode(model.generate(**model_input, max_new_tokens=100, repetition_penalty=1.15)[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30885538-b32e-4911-8a87-8023ecdaaf26",
   "metadata": {
    "id": "VCJnpZoayRgq"
   },
   "source": [
    "### Sweet... it generated content on the human heart! \n",
    "\n",
    "Clearly, this is a bio model. Pretty excited to play around with it, and would love feedback from people in biology about the model's performance on their own tasks!\n",
    "\n",
    "I hope you enjoyed this tutorial on running BioMistral. If you have any questions or suggestions, feel free to reach out to me on [X](https://x.com/harperscarroll) or on the [Discord](https://discord.gg/T9bUNqMS8d).\n",
    "\n",
    "🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
